% AMTH 358 Final - Question 1
% Radix-2, decimation-in-time, FFT
% Definition of fwd DFT: xf(k) = Sum{n=1,N-1} x(k)*exp(-i*2*pi*n*k/N)
% Note that the scaling factor 1/N is omitted here and will be
% applied to the inverse DFT in the same manner as Matlab's native FFT
%--------------------------------------------------------------------------------

function [xf] = fft_dit(x, debug, strict)
  if nargin <3
    strict=1;
  end
  if nargin <2
    debug=1;
  end
  if nargin <1
    usage();
    xf = [];
    return;
  end

  N=length(x);
  L=fix(log2(N));
  xf = NaN * ones(1,N); % output vector is initialized in case of error condition    

  if N ~= 2^L
    debug_disp('Error: input vector x must have size as power of 2',debug);
    return;
  elseif L>10
    if strict
      debug_disp('Error: input vector length must have a size of 1024 or less',debug);
      return;    
    else
      debug_disp('Info: input vector is pretty big!',debug);
    end
  elseif L==0  % must cover this corner case too!
    debug_disp('Info: input vector has length 1', debug);
    xf=x;
    return;
  end

  debug_disp('N roots of unity:',debug);
  W=gen_unit_roots(N);
  debug_disp(W,debug);
  
  % plot_roots(W);  % for debugging
  
  debug_disp('Bit reversed address input sequence',debug);

  xf = bit_rev_sort(x);
  debug_disp(xf,debug);
  
  % construct the radix-2 decimation in time butterfly network
  for stage=0:L-1
    xf_next = [];
    
    k1 = 1;
    W_index = 1;
    
    for bfly=0:N/2 -1
      k2 = k1 + 2^stage;
      
      din = [xf(k1) xf(k2)];
      dout = butterfly(din, W(W_index));

      xf_next(k1) = dout(1);
      xf_next(k2) = dout(2);
      
      if mod(bfly+1,2^stage) == 0
        k1=k1+ (2^stage)+1;
        W_index = 1;
      else
        k1=k1+1;
        W_index = W_index + 2^(L-1-stage);
      end
    end
    xf = xf_next;
  end
  
  
function [dout] = butterfly(din, coefficient)
  din2_scaled = din(2) * coefficient;
  dout(1) = din(1) + din2_scaled;
  dout(2) = din(1) - din2_scaled;
  
function [x_br] = bit_rev_sort(x)
  N=length(x);
  L=fix(log2(N));
  x_br = zeros(1,N);
  for i=1:N
    binary_index_minus1 = dec2bin(i-1,L);
    bit_rev_index = bin2dec(fliplr(binary_index_minus1))+1;
    x_br(bit_rev_index) = x(i);
  end
  
  
function [W] = gen_unit_roots(N)
  % generate a table containing Nth roots of unity
  W = [];
  for k=0:N-1
    root = exp(-i*2*pi*k/N);
    W = [W root];
  end
  
function plot_roots(W)
  % plot the N roots of unity to gain intuition
  figure;
  plot(W,'*');
  grid;
  title('N unity roots');
  xlabel('Real'); ylabel('Imag');
  hold on;
  syms k real;
  kr = linspace(0,2*pi,128);
  c = exp(i*k);
  cr = subs(c,k,kr);
  plot(cr,'-.');
  hold off;

function debug_disp(str, debug)
  if debug == 1
    disp(str);
  end

function usage()
  disp('Usage: fft_dit <data vector> [debug] [strict]');
  disp('               data vector has length equal to a power of two');
  disp('               debug = 1    print out messages to console (default)');
  disp('               strict = 1   enforce limits on vector size (default)');
  
  